package com.mini.gateway.filter;

import cn.hutool.core.util.StrUtil;
import cn.hutool.json.JSONObject;
import cn.hutool.json.JSONUtil;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.GlobalFilter;
import org.springframework.cloud.gateway.filter.factory.rewrite.CachedBodyOutputMessage;
import org.springframework.cloud.gateway.support.BodyInserterContext;
import org.springframework.core.Ordered;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DataBufferUtils;
import org.springframework.http.HttpHeaders;
import org.springframework.http.ReactiveHttpOutputMessage;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpRequestDecorator;
import org.springframework.stereotype.Component;
import org.springframework.web.reactive.function.BodyInserter;
import org.springframework.web.reactive.function.BodyInserters;
import org.springframework.web.reactive.function.server.HandlerStrategies;
import org.springframework.web.reactive.function.server.ServerRequest;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.nio.charset.StandardCharsets;

/**
 * 自定义全局过滤器缓存请求体。
 * 以下代码同样可以使用过滤器工厂实现，不过过滤器工厂需要在配置文件中配置才能生效
 *
 * @author Zhou Feng
 */
@Component
public class CacheRequestBodyFilter implements GlobalFilter, Ordered {
    Logger logger = LoggerFactory.getLogger(CacheRequestBodyFilter.class);

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
        return cacheBody2(exchange, chain);
    }

    //方式1
    private Mono<Void> cacheBody1(ServerWebExchange exchange, GatewayFilterChain chain) {
        return DataBufferUtils.join(exchange.getRequest().getBody())
                .flatMap(dataBuffer -> {
                    if (dataBuffer.readableByteCount() <= 0) {
                        return chain.filter(exchange);
                    }
                    byte[] bytes = new byte[dataBuffer.readableByteCount()];
                    //将请求体中的数据读入 byte[]
                    dataBuffer.read(bytes);
                    //释放缓冲区
                    DataBufferUtils.release(dataBuffer);
                    logger.debug("\n请求体为: {}", new String(bytes, StandardCharsets.UTF_8));

                    //将请求体再次封装到request里，传到下一级，
                    //否则，由于请求体已被消费，后续的服务将取不到值
                    ServerHttpRequest mutatedRequest = new ServerHttpRequestDecorator(
                            exchange.getRequest()) {
                        @Override
                        public Flux<DataBuffer> getBody() {
                            DataBuffer buffer = exchange.getResponse().bufferFactory().wrap(bytes);
                            return Flux.from(Mono.just(buffer));
                        }

                        //一定要重写
                        @Override
                        public HttpHeaders getHeaders() {
                            HttpHeaders httpHeaders = new HttpHeaders();
                            httpHeaders.putAll(super.getHeaders());
                            //由于修改了请求体的body，导致content-length长度不确定，
                            //因此需要删除原先的content-length,
                            //并且设置 Transfer-Encoding: chunked
                            httpHeaders.remove(HttpHeaders.CONTENT_LENGTH);
                            httpHeaders.set(HttpHeaders.TRANSFER_ENCODING, "chunked");
                            return httpHeaders;
                        }
                    };
                    //封装request，传给下一级
                    return chain.filter(exchange.mutate().request(mutatedRequest).build());
                });
    }

    //方式2
    private Mono<Void> cacheBody2(ServerWebExchange exchange, GatewayFilterChain chain) {
        // 读取请求体并修改请求体
        ServerRequest serverRequest = ServerRequest.create(exchange, HandlerStrategies.withDefaults().messageReaders());
        Mono<String> modifiedBody = serverRequest.bodyToMono(String.class)
                .flatMap(body -> {
                    if (StrUtil.isNotBlank(body)) {
                        System.out.println(body);
                        // 对原先的body进行修改操作
                        JSONObject jsonObject = JSONUtil.parseObj(body);
                        jsonObject.set("parentId", 1);
                        return Mono.just(jsonObject.toString());
                    }
                    return Mono.empty();
                });
        BodyInserter<Mono<String>, ReactiveHttpOutputMessage> bodyInserter = BodyInserters.fromPublisher(modifiedBody, String.class);
        HttpHeaders headers = new HttpHeaders();
        CachedBodyOutputMessage outputMessage = new CachedBodyOutputMessage(exchange, headers);
        return bodyInserter.insert(outputMessage, new BodyInserterContext())
                .then(Mono.defer(() -> {
                    ServerHttpRequestDecorator decorator = new ServerHttpRequestDecorator(
                            exchange.getRequest()) {
                        @Override
                        public HttpHeaders getHeaders() {
                            long contentLength = headers.getContentLength();
                            HttpHeaders httpHeaders = new HttpHeaders();
                            httpHeaders.putAll(super.getHeaders());
                            if (contentLength > 0) {
                                httpHeaders.setContentLength(contentLength);
                            } else {
                                httpHeaders.remove(HttpHeaders.CONTENT_LENGTH);
                                httpHeaders.set(HttpHeaders.TRANSFER_ENCODING, "chunked");
                            }
                            return httpHeaders;
                        }

                        //重新封装请求体
                        @Override
                        public Flux<DataBuffer> getBody() {
                            return outputMessage.getBody();
                        }
                    };
                    return chain.filter(exchange.mutate().request(decorator).build());
                }));
    }

    @Override
    public int getOrder() {
        return -1;
    }
}
